import socket
import struct
import json
import http.client
import http.server
import socketserver
import os
import threading
import time

# 多播相关设置
MULTICAST_GROUP_IP = '224.0.0.251'
MULTICAST_PORT = 5353

# HTTP 服务器相关设置
OTA_FILENAME = 'ota.bin'
OTA_FILEPATH = './uploads/ota.bin'
HTTP_SERVER_PORT = 3000

# HTTP 客户端相关设置
TARGET_DEVICE_IP = '192.168.43.1'  # 这里应替换为目标设备的 IP 地址
OTA_TRIGGER_PORT = 8080

# 版本信息
current_version = None
new_version = None

# 多播接收线程
def multicast_listener():
    global current_version, new_version
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(('', MULTICAST_PORT))
    group = socket.inet_aton(MULTICAST_GROUP_IP)
    mreq = struct.pack('4sL', group, socket.INADDR_ANY)
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
    print(f"Listening for multicast messages on {MULTICAST_GROUP_IP}:{MULTICAST_PORT}...")
    
    try:
        while True:
            data, addr = sock.recvfrom(1024)
            print(f"Received message from {addr}:")
            try:
                json_data = json.loads(data.decode('utf-8'))
                ota_state = json_data.get('OTA_state')
                ota_version = json_data.get('OTA_version')
                if ota_state and ota_version:
                    print(f"OTA State: {ota_state}, Version: {ota_version}")
                    # 如果是首次获取版本信息
                    if current_version is None:
                        current_version = ota_version
                        print(f"Initial Version: {current_version}")
                    else:
                        # 检查版本号是否更新
                        new_version = ota_version
                        print(f"New Version: {new_version}")
                        # 比较版本号
                        if new_version > current_version:
                            print("Upgrade successful")
                        else:
                            print("Upgrade failed or version unchanged")
                else:
                    print("Received message is not in the expected format")
            except json.JSONDecodeError:
                print("Received message is not valid JSON")
            print()  # 添加一个空行以便阅读
    except KeyboardInterrupt:
        print("Exiting...")
    finally:
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_DROP_MEMBERSHIP, mreq)
        sock.close()

# HTTP服务器部分
class MyHttpRequestHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/api/uploadfile':
            if not os.path.exists(OTA_FILEPATH):
                self.send_error(404, 'File Not Found')
                return

            file_size = os.path.getsize(OTA_FILEPATH)
            self.send_response(200)
            self.send_header('Content-Disposition', f'attachment; filename="{OTA_FILENAME}"')
            self.send_header('Content-Type', 'application/octet-stream')
            self.send_header('Content-Length', file_size)
            self.end_headers()

            bytes_sent = 0
            with open(OTA_FILEPATH, 'rb') as f:
                while True:
                    chunk = f.read(4096)
                    if not chunk:
                        break
                    bytes_sent += len(chunk)
                    progress = (bytes_sent / file_size) * 100
                    print(f"Sending {len(chunk)} bytes... ({progress:.2f}%)")
                    self.wfile.write(chunk)
            print("File download complete")
        else:
            super().do_GET()

# HTTP客户端部分
def trigger_ota(target_ip):
    try:
        conn = http.client.HTTPConnection(f'{target_ip}:{OTA_TRIGGER_PORT}')
        conn.request('POST', '/ota-start', body='start_ota', headers={'Content-Type': 'text/plain'})
        response = conn.getresponse()
        print(f'Status: {response.status}, Response: {response.read().decode()}')
        conn.close()
    except Exception as e:
        print(f'Error occurred: {e}')

if __name__ == '__main__':
    # 创建 uploads 文件夹
    if not os.path.exists('uploads'):
        os.makedirs('uploads')
    
    # 启动多播接收线程
    multicast_thread = threading.Thread(target=multicast_listener)
    multicast_thread.daemon = True
    multicast_thread.start()
    print("Started multicast listener thread")

    # 启动HTTP服务器
    server_thread = threading.Thread(target=lambda: socketserver.TCPServer(('', HTTP_SERVER_PORT), MyHttpRequestHandler).serve_forever())
    server_thread.daemon = True
    server_thread.start()
    print(f'Server running on port {HTTP_SERVER_PORT}')

    # 等待一会儿让多播线程初始化并接收初始版本信息
    print("Waiting for initial version information...")
    time.sleep(2)  # 等待2秒以确保多播线程有足够时间接收初始版本信息

    # 触发OTA更新
    print("Triggering OTA update...")
    trigger_ota(TARGET_DEVICE_IP)

    # 等待一会儿让设备完成升级
    print("Waiting for device to complete upgrade...")
    time.sleep(10)  # 等待60秒（根据实际升级时间调整）

    # 停止服务器和多播监听（可选）
    try:
        while True:
            pass
    except KeyboardInterrupt:
        print("\nServer shutting down.")
